"""
2.	对下面附件文件（附件.txt）所示的网络输出数据，进行nms（非极大值抑制）处理，
最后用图形显示nms处理前后的效果对比（共20分）
题目要求：
①	显示nms处理前的所有框
②	显示nms处理后的所有框
③	nms算法处理结果正确
④	数据做合适的排序
⑤	处理逻辑正确
"""

import numpy as np
import matplotlib.pyplot as plt


def nms(boxes, plt_params, thresh=0.7):
    M = len(boxes)
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    scores = boxes[:, 4]
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)

    def sep(label=''):
        print('-' * 32, label, '-' * 32)

    def draw_boxes(idx, title):
        plt_params['spn'] += 1
        if plt_params['spn'] > plt_params['spr'] * plt_params['spc']:
            return
        plt.subplot(plt_params['spr'], plt_params['spc'], plt_params['spn'])
        plt.title(title)
        c = 'k'
        plt.plot([x1[idx], x2[idx]], [y1[idx], y1[idx]], color=c)
        plt.plot([x1[idx], x2[idx]], [y2[idx], y2[idx]], color=c)
        plt.plot([x1[idx], x1[idx]], [y1[idx], y2[idx]], color=c)
        plt.plot([x2[idx], x2[idx]], [y1[idx], y2[idx]], color=c)
        plt.ylim(plt.ylim()[::-1])

    draw_boxes(np.arange(M), 'Before')

    order = np.argsort(scores)[::-1]
    idx_arr_keep = np.array([], dtype=np.int32)
    idx_arr_rest = order.copy()
    no = 0
    while True:
        if len(idx_arr_rest) == 0:
            break
        no += 1
        sep(no)
        idx_keep = idx_arr_rest[0]
        idx_arr_keep = np.append(idx_arr_keep, [idx_keep])
        idx_arr_rest = idx_arr_rest[1:]
        print('Before')
        print('To keep:', idx_arr_keep)
        print('Rest:', idx_arr_rest)
        if len(idx_arr_rest) == 0:
            break

        max_x1 = np.maximum(x1[idx_keep], x1[idx_arr_rest])
        min_x2 = np.minimum(x2[idx_keep], x2[idx_arr_rest])
        max_y1 = np.maximum(y1[idx_keep], y1[idx_arr_rest])
        min_y2 = np.minimum(y2[idx_keep], y2[idx_arr_rest])
        xx = min_x2 - max_x1 + 1
        yy = min_y2 - max_y1 + 1
        xx = np.maximum(xx, 0)
        yy = np.maximum(yy, 0)
        I = xx * yy
        U = areas[idx_keep] + areas[idx_arr_rest] - I
        iou = I / (U + 1e-20)
        idx_arr_rest = idx_arr_rest[iou < thresh]

        print('After')
        print('To keep:', idx_arr_keep)
        print('Rest:', idx_arr_rest)

    draw_boxes(idx_arr_keep, 'After')
    return idx_arr_keep


if '__main__' == __name__:
    boxes = np.array([
        [100, 100, 210, 210, 0.72],
        [280, 290, 420, 420, 0.8],
        [220, 220, 320, 330, 0.92],
        [105, 90, 220, 210, 0.71],
        [230, 240, 325, 330, 0.81],
        [305, 300, 420, 420, 0.9],
        [215, 225, 305, 328, 0.6],
        [150, 260, 290, 400, 0.99],
        [102, 108, 208, 208, 0.72],
    ])

    plt_params = {
        'spr': 1,
        'spc': 2,
        'spn': 0
    }
    plt.figure(figsize=[10, 5])
    nms(boxes, plt_params)
    plt.show()
